import pymunk
import pymunk.util as util
from pymunk import Vec2d
import math
import sys
import random

from WorldArm import *
from State import *

class World:
    def __init__(self):
        self.sizeX = 100
        self.sizeY = 100
        self.armLength = 25
        self.armControlLength = 5
        self.armStartPosition = Vec2d(40, 50)
        
        self.k1 = 5000
        self.k2 = 0
        self.damping = 2

        # dx vales for the two controls an arm1
        self.arm1Control1 = 0
        self.arm1Control2 = 0
        
        self.arm1Shape1 = None
        self.arm1body2 = None
        
        self.balls = []
        self.segments = [] # For floors etc
        self.visualBallConnections = []
        
        pymunk.init_pymunk()
        self.space = pymunk.Space(iterations = 1, elastic_iterations=250)
        self.space.gravity = (0.0, -900.0)
    
        self.space.resize_static_hash()
        self.space.resize_active_hash()
        
        # Add fixed body with inf mass and inf moment
        self.fixedBody = pymunk.Body(pymunk.inf, pymunk.inf)
        self.fixedBody.position = self.armStartPosition

        self.arms = []
        self.arms.append(WorldArm(self, self.fixedBody))
    
        self.addFloor()
        self.addWalls()

    def forceOnArm(self, dx):
        if dx < 0:
            return 0
        return self.k1*dx + self.k2*dx*dx
    
    def add_ball(self):
        mass = 1
        radius = 10
        inertia = pymunk.moment_for_circle(mass, 0, radius)
        body = pymunk.Body(mass, inertia)
        body.position = 30, 80
        shape = pymunk.Circle(body, radius)
        shape.elasticity = 0.95
        self.space.add(body, shape)
        self.balls.append(shape)
    
    def addFloor(self):
        body = pymunk.Body(pymunk.inf, pymunk.inf)
        body.position = (1, 5)
        l1 = pymunk.Segment(body, (0, 0), (98.0, 0.0), 1.0)
        l1.elasticity = 0.95
        self.space.add(l1)
        self.segments.append(l1)

    def addWalls(self):
        body = pymunk.Body(pymunk.inf, pymunk.inf)
        body.position = (1, 5)
        l1 = pymunk.Segment(body, (0, 0), (0.0, 99.0), 1.0)
        l1.elasticity = 0.95
        self.space.add(l1)
        self.segments.append(l1)
        body = pymunk.Body(pymunk.inf, pymunk.inf)
        body.position = (99, 5)
        l1 = pymunk.Segment(body, (0, 0), (0.0, 99.0), 1.0)
        l1.elasticity = 0.95
        self.space.add(l1)
        self.segments.append(l1)
    
    def addArm2(self):
        # Body1 is closest to arm1
        body1 = pymunk.Body(10, 100)
        body1.position = (self.armStartPosition.x + self.armLength, self.armStartPosition.y)
        shape = pymunk.Circle(body1, 5)
        shape.elasticity = 0.95
        shape.group = 1
        self.space.add(body1, shape)
        self.balls.append(shape)

        # Body1 is at the end of the arm
        body2 = pymunk.Body(10, 100)
        body2.position = (self.armStartPosition.x + 2*self.armLength, self.armStartPosition.y)    
        shape = pymunk.Circle(body2, 5)
        shape.elasticity = 0.95
        shape.group = 1
        self.space.add(body2, shape)
        self.balls.append(shape)
        self.arm2body2 = body2

        # Constrain the two circles to move together
        upperJoint = pymunk.PinJoint(body2, body1, (0,5), (0,-5))    
        lowerJoint = pymunk.PinJoint(body2, body1, (0,-5), (0,5))    
        self.space.add(upperJoint)
        self.space.add(lowerJoint)

        self.visualBallConnections.append( (len(self.balls) - 1, len(self.balls) - 2) ) 
        
        # Constrain arm1 to arm1
        joint = pymunk.PinJoint(self.arm1body2, body1, (0,0), (0,0))    
        self.space.add(joint)

    def addArm3(self):
        # Body1 is closest to arm1
        body1 = pymunk.Body(10, 100)
        body1.position = (self.armStartPosition.x + 2*self.armLength, self.armStartPosition.y)
        shape = pymunk.Circle(body1, 5)
        shape.elasticity = 0.95
        shape.group = 1
        self.space.add(body1, shape)
        self.balls.append(shape)

        # Body1 is at the end of the arm
        body2 = pymunk.Body(10, 100)
        body2.position = (self.armStartPosition.x + 3*self.armLength, self.armStartPosition.y)    
        shape = pymunk.Circle(body2, 5)
        shape.elasticity = 0.95
        shape.group = 1
        self.space.add(body2, shape)
        self.balls.append(shape)
        body2.group = 1    

        # Constrain the two circles to move together
        upperJoint = pymunk.PinJoint(body2, body1, (0,5), (0,-5))    
        lowerJoint = pymunk.PinJoint(body2, body1, (0,-5), (0,5))    
        self.space.add(upperJoint)
        self.space.add(lowerJoint)

        self.visualBallConnections.append( (len(self.balls) - 1, len(self.balls) - 2) ) 
        
        # Constrain arm1 to arm1
        joint = pymunk.PinJoint(self.arm2body2, body1, (0,0), (0,0))    
        self.space.add(joint)

    def setForcesFromDamping(self):
        if self.arm1Shape2 != None:
            body = self.arm1Shape2.body
            body.reset_forces()
            v = body.velocity
            body.apply_force( -self.damping*v, (0,0) )
        
    def setForcesFromControls(self):
        if self.arm1Shape1 == None:
            return
        body = self.arm1Shape1.body
        body.reset_forces()
        # Upper control
        dx = self.arm1Control1 - math.sin(body.angle) * self.armControlLength
        f = -self.forceOnArm(dx)
        if f<0:
            fv = (f, 0)
            p = (0, math.cos(body.angle) * self.armControlLength)
            #print "1 :", body.angle, dx, fv, p
            body.apply_force(fv, p)
        # Lower control
        dx = self.arm1Control2 + math.sin(body.angle) * self.armControlLength
        f = -self.forceOnArm(dx)
        if f<0:
            fv = (f, 0)
            p = (0, -math.cos(body.angle) * self.armControlLength)
            #print "1 :", body.angle, dx, fv, p
            body.apply_force(fv, p)
    
    def setForceFromStops(self):
        if self.arm1Shape1 == None:
            return
        body = self.arm1Shape1.body
        dx = self.armControlLength*math.cos(body.angle)
        body2 = self.arm1Shape2.body
        if dx < 1.0 and body2.velocity.x < 0:
            v = body2.velocity
            body2.apply_force( (-100*v.x,0), (0,0) )        
    
    def implement(self, action):
        c1 = action[0]
        c2 = action[1]
        if c1 < -1:
            c1 = -1
        elif c1 > 10:
            c1 = 10
        if c2 < -1:
            c2 = -1
        elif c2 > 10:
            c2 = 10
        print "World: Implementing {0} {1} ".format(c1, c2)
        self.arm1Control1 = c1
        self.arm1Control2 = c2
            
    def getState(self):
        state = State()
        state.arm1.c1 = self.arm1Control1
        state.arm1.c2 = self.arm1Control2
        state.arm1.a = self.arm1Shape1.body.angle
        state.arm1.aDot = self.arm1Shape1.body.angular_velocity
        return state
    
    def step(self, dt):
        N = 10
        for i in range(N):
            self.setForcesFromControls()
            self.setForcesFromDamping()
            self.setForceFromStops()
            self.space.step(dt/N)
        